"""
@author: Viet Nguyen <nhviet1009@gmail.com>
"""
import argparse

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageOps
from utils import get_data


def get_args():
    parser = argparse.ArgumentParser("Image to ASCII")
    parser.add_argument("--input", type=str, default="data/input.jpg", help="Path to input image")
    parser.add_argument("--output", type=str, default="data/output.jpg", help="Path to output text file")
    parser.add_argument("--language", type=str, default="english")
    parser.add_argument("--mode", type=str, default="standard")
    parser.add_argument("--background", type=str, default="black", choices=["black", "white"],
                        help="background's color")
    parser.add_argument("--num_cols", type=int, default=300, help="number of character for output's width")
    parser.add_argument("--scale", type=int, default=2, help="upsize output")
    args = parser.parse_args()
    return args


def main(opt):
    if opt.background == "white":
        背景编码 = (255, 255, 255)
    else:
        背景编码 = (0, 0, 0)
    字符列表, font, 字符示例, scale = get_data(opt.language, opt.mode)
    字符种数 = len(字符列表)
    列数 = opt.num_cols
    图片 = cv2.imread(opt.input, cv2.IMREAD_COLOR)
    图片 = cv2.cvtColor(图片, cv2.COLOR_BGR2RGB)
    图高, 图宽, _ = 图片.shape

    字符宽度, 字符高度 = font.getsize(字符示例)
    字符高宽比 = 字符高度/字符宽度

    格宽 = 图宽 / opt.num_cols
    格高 = 字符高宽比 * 格宽
    行数 = int(图高 / 格高)
    if 列数 > 图宽 or 行数 > 图高:
        print("Too many columns or rows. Use default setting")
        格宽 = 6
        格高 = 字符高宽比 * 格宽
        列数 = int(图宽 / 格宽)
        行数 = int(图高 / 格高)
    输出宽度 = 字符宽度 * 列数
    输出高度 = scale * 字符高度 * 行数
    输出图片 = Image.new("RGB", (输出宽度, 输出高度), 背景编码)
    draw = ImageDraw.Draw(输出图片)
    for 行号 in range(行数):
        for 列号 in range(列数):
            局部图 = 图片[int(行号 * 格高) : min(int((行号 + 1) * 格高), 图高),
                        int(列号 * 格宽) : min(int((列号 + 1) * 格宽), 图宽)]
            局部平均颜色 = np.sum(np.sum(局部图, axis=0), axis=0) / (格高 * 格宽)
            局部平均颜色 = tuple(局部平均颜色.astype(np.int32).tolist())
            字符 = 字符列表[min(int(np.mean(局部图) * 字符种数 / 255), 字符种数 - 1)]
            draw.text((列号 * 字符宽度, 行号 * 字符高度), 字符, fill=局部平均颜色, font=font)

    if opt.background == "white":
        cropped_image = ImageOps.invert(输出图片).getbbox()
    else:
        cropped_image = 输出图片.getbbox()
    输出图片 = 输出图片.crop(cropped_image)
    输出图片.save(opt.output)


if __name__ == '__main__':
    opt = get_args()
    main(opt)
